import json

import numpy as np
import pandas as pd


def get_label_dictionary(labels, keys):
    """Associate key (filename) to value (box coords, class)"""
    dictionary = {}
    for key in keys:
        dictionary[key] = []  # empty boxes

    for label in labels:
        if len(label) != 6:
            print("Incomplete label:", label[0])
            continue

        value = label[1:]

        if value[0] == value[1]:
            continue
        if value[2] == value[3]:
            continue

        if label[-1] == 0:
            print("No object labelled as bg:", label[0])
            continue

        # box coords are float32
        value = value.astype(np.float32)
        new_value = value.copy()
        # x_min,y_min,w,h=>x_min,x_max,y_min,y_max
        new_value[1] = value[2] + value[0]
        new_value[2] = value[1]
        new_value[3] = value[3] + value[1]
        # filename is key
        key = label[0]
        # boxes = bounding box coords and class label
        boxes = dictionary[key]
        boxes.append(new_value)
        dictionary[key] = boxes

    # remove dataset entries w/o labels
    for key in keys:
        if len(dictionary[key]) == 0:
            del dictionary[key]

    return dictionary


def build_label_dictionary(path):
    """Build a dict with key=filename, value=[box coords, class]"""
    train_csv = pd.read_csv(path)
    train_csv = train_csv.query('annotations!="[]"')

    target = []
    for i in range(train_csv.shape[0]):
        item = train_csv.iloc[i]
        annotation = item['annotations']
        annotation = annotation.replace('\'', '\"')
        anno_list = json.loads(annotation)
        image_id = item['image_id']

        for anno in anno_list:
            target.append(
                [image_id, anno['x'], anno['y'], anno['width'],
                 anno['height'], 1])  # 最后一个表示类别
    labels = np.array(target)

    keys = np.unique(labels[:, 0])
    dictionary = get_label_dictionary(labels, keys)
    classes = np.unique(labels[:, -1]).astype(int).tolist()
    # insert background label 0
    classes.insert(0, 0)
    print("Num of unique classes: ", classes)
    return dictionary, classes
